import tensorflow as tf
from python_ai.common.xcommon import *

print(tf.__version__)


@tf.function
def get_y_dydx_dydx2(x):
    a = tf.constant(3.0, name='a')
    b = tf.constant(4.0, name='b')
    c = tf.constant(5.0, name='c')

    with tf.GradientTape() as tape1:
        with tf.GradientTape() as tape2:
            tape2.watch(x)
            y = a * x ** 2 + b * x + c
        tape1.watch(x)
        dydx = tape2.gradient(y, x)
    dydx2 = tape1.gradient(dydx, x)

    return y, dydx, dydx2


sep('x=1.0')
print(get_y_dydx_dydx2(tf.Variable(1.0, name='x1')))
sep('x=2.0')
print(get_y_dydx_dydx2(tf.constant(2.0, name='c1')))
